package com.nanohadoop.yarn;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;

/**
 * YARN集群管理类，提供简单的集群管理接口
 */
public class YarnCluster {
    private static final Logger log = LoggerFactory.getLogger(YarnCluster.class);

    private final ResourceScheduler scheduler;
    private final ResourceManager resourceManager;

    /**
     * 创建YARN集群，使用默认调度器
     */
    public YarnCluster() {
        this.scheduler = new FIFOResourceScheduler(); // 默认使用FIFO调度器
        this.resourceManager = new ResourceManager(this.scheduler);
    }

    /**
     * 创建YARN集群，指定调度器策略
     */
    public YarnCluster(ResourceScheduler scheduler) {
        this.scheduler = scheduler;
        this.resourceManager = new ResourceManager(scheduler);
    }

    /**
     * 启动YARN集群
     */
    public void start() {
        try {
            scheduler.start();
            resourceManager.start();
            log.debug("YARN cluster started");
        } catch (IOException e) {
            throw new RuntimeException("Failed to start ResourceManager", e);
        }
    }

    /**
     * 关闭YARN集群
     */
    public void stop() {
        log.debug("Starting YarnCluster.stop()");
        // 停止所有NodeManager
        log.debug("Stopping all NodeManagers");
        for (NodeManager nodeManager : scheduler.getNodeManagers().values()) {
            log.debug("Stopping NodeManager: " + nodeManager.getNodeId());
            nodeManager.stop();
        }
        // 停止ResourceManager
        log.debug("Stopping ResourceManager");
        resourceManager.stop();
        // 停止ResourceScheduler
        log.debug("Stopping ResourceScheduler");
        scheduler.stop();
        log.debug("YARN cluster stopped");
    }

    /**
     * 添加节点到集群
     */
    public void addNode(NodeManager nodeManager) {
        nodeManager.start();
        scheduler.registerNodeManager(nodeManager);
    }

    /**
     * 添加节点到集群
     */
    public void addNode(String nodeId, int memoryMB, int cpuCores) {
        NodeManager nodeManager = new NodeManager(nodeId, memoryMB, cpuCores);
        nodeManager.start();
        scheduler.registerNodeManager(nodeManager);
    }

    /**
     * 提交应用程序到YARN集群
     */
    public String submitApplication(String name, String user, int memoryMBPerContainer, int cpuCoresPerContainer, int numContainers) {
        Resource resource = new Resource(memoryMBPerContainer, cpuCoresPerContainer);
        return scheduler.submitApplication(name, user, resource, numContainers);
    }

    /**
     * 提交应用程序到YARN集群
     */
    public String submitApplication(String appName, String user, Resource resource, int numContainers) {
        return scheduler.submitApplication(appName, user, resource, numContainers);
    }

    /**
     * 完成应用程序
     */
    public void completeApplication(String applicationId) {
        scheduler.completeApplication(applicationId);
    }

    /**
     * 获取资源调度器
     */
    public ResourceScheduler getScheduler() {
        return scheduler;
    }

    /**
     * 获取资源管理器
     */
    public ResourceManager getResourceManager() {
        return resourceManager;
    }
}